#!/usr/bin/env python
import fnmatch
import os
from argparse import ArgumentParser
import numpy as np
from astropy.io import fits
from pylab import *


def raw2fits(fname, row=9222, col=8976, overwrite=False, suffix='raw', dtype='uint16'):
    print('*** processing: {}'.format(fname))
    if fname.endswith(suffix) is False:
        print('Error: suffix does not match')
        return
    fout = fname.replace('.{}'.format(suffix), '.fits')
    if os.path.exists(fout) and overwrite is False:
        print('Warning: fits file already exists'.format(os.path.basename(fout)))
        return

    # data = np.fromfile(fname, dtype=np.dtype('>u2'))
    data = np.fromfile(fname, dtype=dtype)
    #data = fromfile(fname,dtype='<i16')
    data = data.reshape(row, col)
    hdu = fits.PrimaryHDU(data)
    hdu.writeto(fout, overwrite=overwrite)
    print('Success: fits file written')


def main(path, **kwargs):
    if os.path.isfile(path):
        raw2fits(path, **kwargs)

    elif os.path.isdir(path):
        if path.endswith('/'):
            path = path[:-1]
        pattern = '*.{}'.format(kwargs['suffix'])
        flist = []
        for r, d, f in os.walk(path):
            names = fnmatch.filter(f, pattern)
            if len(names) > 0:
                for name in names:
                    flist.append('{}/{}'.format(r, name))
        for f in flist:
            raw2fits(f, **kwargs)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('path', type=str)
    parser.add_argument('--row', '-y', default=9400, type=int)
    parser.add_argument('--col', '-x', default=10000, type=int)
    parser.add_argument('--overwrite', '-o', action='store_true')
    parser.add_argument('--suffix', '-s', default='raw', type=str)
    parser.add_argument('--dtype', '-t', default='int32', type=str)
    args = parser.parse_args()
    main(args.path, row=args.row, col=args.col, overwrite=args.overwrite, suffix=args.suffix, dtype=args.dtype)
